#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 

template<typename TYPE_X1, typename TYPE_CONDITION> class KernelSelectV2 {
    public:
        __aicore__ inline KernelSelectV2() {}
        __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,int32_t M, int32_t N, int32_t Z,
            uint32_t totalLength, uint32_t ALIGN_NUM, uint32_t block_size, uint32_t core_size, uint32_t core_remain) {
            ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
            this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
            this->tileLength = block_size;
            this->ALIGN_NUM = ALIGN_NUM;
            this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);
            this->totalLength = totalLength;
            auto startPointer = core_size * GetBlockIdx();
            auto bufferlength = this->blockLength;

            this->M = M; //行
            this->N = N; //列
            this->Z = Z; //循环次数
    
            Gm_x1.SetGlobalBuffer((__gm__ TYPE_X1*)x1 + startPointer ,bufferlength);
            Gm_x2.SetGlobalBuffer((__gm__ TYPE_X1*)x2 + startPointer , bufferlength);
            Gm_cdt.SetGlobalBuffer((__gm__  uint8_t *)condition + startPointer, bufferlength);
            Gm_y.SetGlobalBuffer((__gm__ TYPE_X1*)y + startPointer , bufferlength);
            
    
            this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);
    
            pipe.InitBuffer(Q_x1, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));
            pipe.InitBuffer(Q_x2, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));
            pipe.InitBuffer(Q_cdt, BUFFER_NUM, this->tileLength * sizeof(uint8_t));
            pipe.InitBuffer(Q_y, BUFFER_NUM, this->tileLength * sizeof(TYPE_X1));

            pipe.InitBuffer(B_bits, this->tileLength * sizeof(uint8_t));

            pipe.InitBuffer(B_half, this->tileLength * sizeof(half));
            //取一个全为0的tensor作为处理
            pipe.InitBuffer(B_zero, this->tileLength * sizeof(half));
            this->zero = B_zero.Get<half>();

            // for(int32_t i=0;i<this->tileLength;i++){

            //     this->zero.SetValue(i,half(0));
            // }
            Duplicate(this->zero, half(0), this->tileLength);

            // for(int32_t i=0;i<bufferlength;i++){
             
            //     Gm_y.SetValue(i,static_cast<uint8_t>(Gm_cdt.GetValue(i)));
            //  }

            


            if constexpr (std::is_same_v<TYPE_X1, int32_t>) {
                pipe.InitBuffer(B_x1, this->tileLength * sizeof(float));
                pipe.InitBuffer(B_x2, this->tileLength * sizeof(float));
                pipe.InitBuffer(B_y, this->tileLength * sizeof(float));
            }
            else if constexpr (std::is_same_v<TYPE_X1, int8_t>) {
                pipe.InitBuffer(B_x1, this->tileLength * sizeof(half));
                pipe.InitBuffer(B_x2, this->tileLength * sizeof(half));
                pipe.InitBuffer(B_y, this->tileLength * sizeof(half));
            }
            
        }
        __aicore__ inline void Process() { //假设每次处理一行来看


            // if ( this->blockLength < 150528){

            //     for(int32_t i=0;i<M * N * Z;i++)
            //     {
                    
            //         TYPE_X1 pos = Gm_cdt.GetValue(i) ? Gm_x1.GetValue(i) : Gm_x2.GetValue(i);
            //         Gm_y.SetValue(i,pos);
            //     }
    


            // }


            // else {

            

            int32_t loopCount = this->tileNum;
            for (int32_t i = 0; i < loopCount-1; i++) {
                CopyIn(i, this->tileLength);
                Compute(i, this->tileLength);
                CopyOut(i, this->tileLength);
            }
            uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
            CopyIn(loopCount - 1, (length +31 ) / 32 * 32);
            Compute(loopCount - 1, (length +31 ) / 32 * 32);
            CopyOut(loopCount - 1, (length +31 ) / 32 * 32);
            //CopyOut(loopCount - 1, length);


            // for(int32_t i=0;i<M * N * Z;i++)
            // {
                
            //     TYPE_X1 pos = Gm_cdt.GetValue(i) ? Gm_x1.GetValue(i) : Gm_x2.GetValue(i);
            //     Gm_y.SetValue(i,pos);
            // }

        //     for(int32_t z1 = 0; z1 < Z; z1++){

        //         for(int32_t i=0;i<M;i++){

        //             for(int32_t j=0;j<N;j++){        

        //                 if(Gm_cdt.GetValue((z1*M*N+i*M+j))==false){

        //                     Gm_y.SetValue(z1*M*N+i*M+j,Gm_x2.GetValue(z1*M*N+i*M+j));

        //                 }

        //                 // else{

        //                 //     Gm_y.SetValue(z1*M*N+i*M+j,Gm_x1.GetValue(z1*M*N+i*M+j));

        //                 // }
        //             }
        //         }

        //     }
       // }

       
           
        }   
    
         
    private:
        __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
           LocalTensor<TYPE_X1> x1 = Q_x1.AllocTensor<TYPE_X1>();
           LocalTensor<TYPE_X1> x2 = Q_x2.AllocTensor<TYPE_X1>();
           LocalTensor<uint8_t> cdt = Q_cdt.AllocTensor<uint8_t>();
           DataCopy(x1, Gm_x1[progress * this->tileLength], length);
           DataCopy(x2, Gm_x2[progress * this->tileLength], length);
           
           DataCopy(cdt, (Gm_cdt[progress * this->tileLength]), length);
           
           Q_x1.EnQue(x1);
           Q_x2.EnQue(x2);
           Q_cdt.EnQue(cdt);
        }
        __aicore__ inline void Compute(int32_t progress, uint32_t length) {
            LocalTensor<TYPE_X1> x1 = Q_x1.DeQue<TYPE_X1>();
            LocalTensor<TYPE_X1> x2 = Q_x2.DeQue<TYPE_X1>();
            LocalTensor<uint8_t> cdt = Q_cdt.DeQue<uint8_t>();
            LocalTensor<TYPE_X1> y = Q_y.AllocTensor<TYPE_X1>();

            auto bits = B_bits.Get<uint8_t>();
            if constexpr (std::is_same_v<TYPE_X1, half>) {

                auto p1 = B_half.Get<half>();

                Cast(p1, cdt, RoundMode::CAST_NONE, length);

                Compare(bits, p1,zero,CMPMODE::NE,length);
                Select(y, bits, x1, x2, SELMODE::VSEL_TENSOR_TENSOR_MODE, length);
            }
            else if constexpr (std::is_same_v<TYPE_X1, float>) {

                auto p1 = B_half.Get<half>();

                Cast(p1, cdt, RoundMode::CAST_NONE, length);

                Compare(bits, p1,zero,CMPMODE::NE,length);

                Select(y, bits, x1, x2, SELMODE::VSEL_TENSOR_TENSOR_MODE, length);
            }
            else if constexpr (std::is_same_v<TYPE_X1, int8_t>) {

                auto p1 = B_half.Get<half>();

                Cast(p1, cdt, RoundMode::CAST_NONE, length);

                Compare(bits, p1,zero,CMPMODE::NE,length);
                auto half_x1 = B_x1.Get<half>();
                auto half_x2 = B_x2.Get<half>();
                auto half_y = B_y.Get<half>();
                Cast(half_x1, x1, RoundMode::CAST_NONE, length);
                Cast(half_x2, x2, RoundMode::CAST_NONE, length);
                //Cast(half_y, y, RoundMode::CAST_NONE, length);
                Select(half_y, bits, half_x1, half_x2, SELMODE::VSEL_TENSOR_TENSOR_MODE, length);
                Cast(y, half_y, RoundMode::CAST_NONE, length);
            }
            else{

                auto p1 = B_half.Get<half>();

                Cast(p1, cdt, RoundMode::CAST_NONE, length);

                Compare(bits, p1,zero,CMPMODE::NE,length);

                auto float_x1 = B_x1.Get<float>();
                auto float_x2 = B_x2.Get<float>();
                auto float_y = B_y.Get<float>();
                Cast(float_x1, x1, RoundMode::CAST_NONE, length);
                Cast(float_x2, x2, RoundMode::CAST_NONE, length);
                //Cast(float_y, y, RoundMode::CAST_RINT, length);
                Select(float_y, bits, float_x1, float_x2, SELMODE::VSEL_TENSOR_TENSOR_MODE, length);
                Cast(y, float_y, RoundMode::CAST_RINT, length);
            }
            Q_x1.FreeTensor(x1);  
            Q_x2.FreeTensor(x2);
            Q_cdt.FreeTensor(cdt);
            Q_y.EnQue<TYPE_X1>(y);
        }
        __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
            LocalTensor<TYPE_X1> y = Q_y.DeQue<TYPE_X1>();
            DataCopy(Gm_y[progress * this->tileLength], y, length);
            Q_y.FreeTensor(y);
        }
    private:
        TPipe pipe;
        TQue<QuePosition::VECIN, BUFFER_NUM> Q_x1, Q_x2, Q_cdt;
        TQue<QuePosition::VECOUT, BUFFER_NUM> Q_y;

        TBuf<QuePosition::VECCALC> B_result, B_zero, B_bits,B_half;
        //TBuf<QuePosition::VECCALC> B_x1, B_x2;
        TBuf<QuePosition::VECCALC> B_x1, B_x2, B_y;

        int32_t M;
        int32_t N;
        int32_t Z;

        LocalTensor<half> zero;
      
        GlobalTensor<TYPE_X1> Gm_x1, Gm_x2, Gm_y;
        GlobalTensor<uint8_t> Gm_cdt;
        uint32_t blockLength;
        uint32_t tileNum;
        uint32_t totalLength;
        uint32_t tileLength;
        uint32_t ALIGN_NUM;
    };
    
   

template<typename TYPE_X1, typename TYPE_CONDITION> class KernelSelectV2_Broadcast {
    using T = TYPE_X1;
public:
    __aicore__ inline KernelSelectV2_Broadcast() {}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,int32_t M, int32_t N, int32_t Z,
                                int32_t y_dimensional, 
                                int32_t* y_ndarray, int32_t* cdt_ndarray, int32_t* x1_ndarray, int32_t* x2_ndarray,
                                int32_t* y_sumndarray, int32_t* cdt_sumndarray, int32_t* x1_sumndarray, int32_t* x2_sumndarray) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        // this->tileDataNum = tileDataNum;
        // this->coreDataNum = CoreDataNum;
        // this->tileNum = finalTileNum;
        // this->tailDataNum = TailDataNum;

        
        this->M = M; //行
        this->N = N; //列
        this->Z = Z; //循环次数

        this->y_dimensional = y_dimensional;

        this->y_ndarray = y_ndarray;
        this->cdt_ndarray = cdt_ndarray;
        this->x1_ndarray = x1_ndarray;
        this->x2_ndarray = x2_ndarray;

        this->y_sumndarray = y_sumndarray;
        this->cdt_sumndarray = cdt_sumndarray;
        this->x1_sumndarray = x1_sumndarray;
        this->x2_sumndarray = x2_sumndarray;

        Gm_x1.SetGlobalBuffer((__gm__ TYPE_X1*)x1 , 1);
        Gm_x2.SetGlobalBuffer((__gm__ TYPE_X1*)x2 , 1);
        Gm_cdt.SetGlobalBuffer((__gm__ TYPE_CONDITION*)condition , 1);
        Gm_y.SetGlobalBuffer((__gm__ TYPE_X1*)y , 1);
    }
    __aicore__ inline void Process() {

        int dim = this->y_dimensional;
        
        for(int j=0; j<this->y_sumndarray[dim]; j++)
        {
            int cdt_start = 0, x1_start = 0, x2_start = 0;
            for(int k=0; k<dim; k++)
            {
                if(this->cdt_ndarray[k] != 1){
                    cdt_start += this->cdt_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
                if(this->x1_ndarray[k] != 1){
                    x1_start += this->x1_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
                if(this->x2_ndarray[k] != 1){
                    x2_start += this->x2_sumndarray[k] * (j / this->y_sumndarray[k] % this->y_ndarray[k]);
                }
            }


            TYPE_X1 pos = Gm_cdt.GetValue(cdt_start) ? Gm_x1.GetValue(x1_start) : Gm_x2.GetValue(x2_start);
            Gm_y.SetValue(j,pos);
            // float start = startGm.GetValue(start_start);
            // float end = endGm.GetValue(end_start);
            // float weight = weightGm.GetValue(weight_start);
            // float y;
            // y = start + weight * (end - start);
            // yGm.SetValue(j, (DTYPE_Y)y);


        }
    } 
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueSTART, inQueueEND, inQueueWEIGHT;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3,boox;

    int32_t M;
    int32_t N;
    int32_t Z;

    
      
    GlobalTensor<TYPE_X1> Gm_x1, Gm_x2, Gm_y;
    GlobalTensor<TYPE_CONDITION> Gm_cdt;

    int32_t y_dimensional;
    int32_t *y_ndarray;
    int32_t *cdt_ndarray;
    int32_t *x1_ndarray;
    int32_t *x2_ndarray;

    int32_t * y_sumndarray;
    int32_t *cdt_sumndarray;
    int32_t *x1_sumndarray;
    int32_t *x2_sumndarray;
};

    extern "C" __global__ __aicore__ void select_v2(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
        GET_TILING_DATA(tiling_data, tiling);
        // TODO: user kernel impl

        if(TILING_KEY_IS(1)){

            KernelSelectV2<DTYPE_X1, DTYPE_CONDITION> op;
            op.Init(condition, x1, x2, y, tiling_data.M, tiling_data.N, tiling_data.Z,
                tiling_data.totalLength, tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain);
            op.Process();

        }

        else if(TILING_KEY_IS(2)){

            KernelSelectV2_Broadcast<DTYPE_X1, DTYPE_CONDITION> op;

            op.Init(condition, x1, x2, y, tiling_data.M, tiling_data.N, tiling_data.Z,
                tiling_data.y_dimensional,
                tiling_data.y_ndarray, tiling_data.cdt_ndarray, tiling_data.x1_ndarray, tiling_data.x2_ndarray,
                tiling_data.y_sumndarray, tiling_data.cdt_sumndarray, tiling_data.x1_sumndarray, tiling_data.x2_sumndarray);  
            op.Process();
        }
    }